"""
Service layer that orchestrates point cloud registration and result loading.

Responsibilities:
    * Reading registration metadata.
    * Executing registration jobs and persisting the resulting transforms.
    * Loading registered groups with colourised error information for the UI.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, List, Optional

import numpy as np
import open3d as o3d
from matplotlib import cm

from config.settings import settings
from core.registration import registration_point_clouds


@dataclass
class RegisteredGroup:
    """Data transfer object representing a registered point cloud group."""

    name: str
    ply: o3d.geometry.PointCloud
    mesh: o3d.geometry.TriangleMesh
    distances: np.ndarray
    threshold: float
    transform: np.ndarray


class RegistrationManager:
    """High level manager for registration workflows and cached results."""

    def __init__(self) -> None:
        self._settings = settings
        self._cache_path = self._settings.paths.registration_cache_file

    # ------------------------------------------------------------------ #
    # Metadata and cache helpers
    # ------------------------------------------------------------------ #
    def load_metadata(self, metadata_path: Optional[Path | str] = None) -> dict:
        metadata_file = Path(metadata_path) if metadata_path else self._settings.paths.metadata_file
        if not metadata_file.exists():
            raise FileNotFoundError(f"未找到元数据文件: {metadata_file}")

        with metadata_file.open("r", encoding="utf-8") as f:
            return json.load(f)

    def _load_cache(self) -> dict:
        if self._cache_path.exists():
            with self._cache_path.open("r", encoding="utf-8") as f:
                return json.load(f)
        return {}

    def _write_cache(self, cache: dict) -> None:
        self._cache_path.parent.mkdir(parents=True, exist_ok=True)
        with self._cache_path.open("w", encoding="utf-8") as f:
            json.dump(cache, f, indent=4)

    def save_registration_result(self, group_idx: int, transformation: np.ndarray) -> None:
        cache = self._load_cache()
        cache[str(group_idx)] = transformation.tolist()
        self._write_cache(cache)

    # ------------------------------------------------------------------ #
    # Processing pipeline
    # ------------------------------------------------------------------ #
    def process_groups(
        self,
        metadata_path: Optional[Path | str] = None,
        output_dir: Optional[Path | str] = None,
    ) -> List[Path]:
        """
        Execute registration for all groups defined in the metadata.

        Returns:
            List of paths to the newly written registered PLY files.
        """
        metadata = self.load_metadata(metadata_path)
        groups = metadata.get("groups", [])

        if not groups:
            return []

        target_dir = Path(output_dir) if output_dir else self._settings.paths.registered
        target_dir.mkdir(parents=True, exist_ok=True)

        generated_paths: List[Path] = []

        for idx, group in enumerate(groups):
            obj_path = target_dir / group["output_file"]
            ply_path = obj_path.with_suffix(".ply")

            if not obj_path.exists() or not ply_path.exists():
                print(f"跳过 group {idx}: 缺少 obj 或 ply 文件 ({obj_path}, {ply_path})")
                continue

            translate = group.get("transform", {}).get("translate", [0, 0, 0])

            source_pcd = o3d.io.read_point_cloud(str(ply_path))
            transform = registration_point_clouds(
                scan_filename=str(ply_path),
                bim_filename=str(obj_path),
            )

            source_pcd.transform(transform)
            source_pcd.translate(translate)

            self.save_registration_result(idx, np.asarray(transform))

            registered_path = target_dir / f"group_{idx + 1}_registered.ply"
            o3d.io.write_point_cloud(str(registered_path), source_pcd)
            generated_paths.append(registered_path)
            print(f"已保存配准结果：{registered_path}")

        return generated_paths

    # ------------------------------------------------------------------ #
    # Data loading for UI
    # ------------------------------------------------------------------ #
    def load_registered_groups(
        self,
        metadata_path: Optional[Path | str] = None,
        apply_cached_transform: bool = True,
        group_indices: Optional[Iterable[int]] = None,
    ) -> List[RegisteredGroup]:
        """
        Load registered groups along with colourised distance metrics.
        """
        metadata = self.load_metadata(metadata_path)
        groups_meta = metadata.get("groups", [])
        cache = self._load_cache() if apply_cached_transform else {}

        registered_groups: List[RegisteredGroup] = []
        selected_indices = set(group_indices) if group_indices is not None else None

        for idx, group in enumerate(groups_meta):
            if selected_indices is not None and idx not in selected_indices:
                continue

            obj_path = self._settings.paths.registered / group["output_file"]
            ply_path = obj_path.with_suffix(".ply")

            if not obj_path.exists() or not ply_path.exists():
                print(f"跳过 group {idx}: 缺少 obj 或 ply 文件 ({obj_path}, {ply_path})")
                continue

            ply = o3d.io.read_point_cloud(str(ply_path))
            mesh = o3d.io.read_triangle_mesh(str(obj_path))

            if not ply.has_points() or not mesh.has_vertices():
                print(f"跳过 group {idx}: 点云或网格为空")
                continue

            mesh.compute_vertex_normals()

            transform_matrix = np.eye(4)
            if str(idx) in cache:
                transform_matrix = np.array(cache[str(idx)])
                ply.transform(transform_matrix)

            translate = group.get("transform", {}).get("translate", [0, 0, 0])
            ply.translate(translate)
            mesh.translate(translate)

            distances, threshold = self._compute_distances_with_colors(ply, mesh)

            registered_groups.append(
                RegisteredGroup(
                    name=f"Group {idx}",
                    ply=ply,
                    mesh=mesh,
                    distances=distances,
                    threshold=threshold,
                    transform=transform_matrix,
                )
            )

        return registered_groups

    def _compute_distances_with_colors(
        self,
        ply: o3d.geometry.PointCloud,
        mesh: o3d.geometry.TriangleMesh,
    ) -> tuple[np.ndarray, float]:
        """
        Compute point-to-surface distances and assign colours to the ply object.
        """
        sample_points = settings.registration.viewer_sample_points
        threshold = settings.registration.default_threshold

        obj_point_cloud = mesh.sample_points_uniformly(number_of_points=sample_points)
        distances = np.array(ply.compute_point_cloud_distance(obj_point_cloud))

        colors = np.full((len(ply.points), 3), 0.5)
        above_threshold = distances > threshold

        if np.any(above_threshold):
            dist_subset = distances[above_threshold]
            denom = max(dist_subset.max() - threshold, 1e-9)
            normalized = (dist_subset - threshold) / denom
            cmap = cm.get_cmap("coolwarm")
            colors[above_threshold] = cmap(normalized)[:, :3]

        ply.colors = o3d.utility.Vector3dVector(colors)
        return distances, threshold


